#!/usr/bin/env python3
"""
Nginx Configuration Manager
A Python script that exposes an API to dynamically update nginx configuration for code-server instances.
"""

import json
import os
import subprocess
import sys
import time
from flask import Flask, request, jsonify
from typing import List, Dict, Any
import logging
import paramiko

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Configuration
NGINX_SITES_AVAILABLE = "/etc/nginx/sites-available"
NGINX_SITES_ENABLED = "/etc/nginx/sites-enabled"
CONFIG_FILE_NAME = "code-server-prod.conf"

class NginxConfigManager:
    def __init__(self):
        self.config_path = os.path.join(NGINX_SITES_AVAILABLE, CONFIG_FILE_NAME)
        self.enabled_path = os.path.join(NGINX_SITES_ENABLED, CONFIG_FILE_NAME)
    
    def _validate_json_data(self, data: List[Dict[str, Any]]) -> bool:
        """Validate the input JSON data"""
        if not isinstance(data, list):
            return False
        
        for item in data:
            if not isinstance(item, dict):
                return False
            
            required_keys = ['uuid', 'ip', 'port']
            if not all(key in item for key in required_keys):
                return False
            
            # Validate data types and basic format
            if not isinstance(item['uuid'], str) or not item['uuid'].strip():
                return False
            if not isinstance(item['ip'], str) or not item['ip'].strip():
                return False
            if not isinstance(item['port'], (int, str)):
                return False
            
            # Try to convert port to int to validate
            try:
                port = int(item['port'])
                if port <= 0 or port > 65535:
                    return False
            except ValueError:
                return False
        
        return True
    
    def _generate_nginx_config(self, instances: List[Dict[str, Any]]) -> str:
        """Generate nginx configuration based on instances data"""
        config_lines = []
        
        # Add server block start
        config_lines.append("server {")
        config_lines.append("    listen 40001;")
        config_lines.append("    listen [::]:40001;")
        config_lines.append("    server_name 192.168.10.231;")
        config_lines.append("")
        
        # Add location blocks for each instance
        for instance in instances:
            uuid = instance['uuid']
            ip = instance['ip']
            port = instance['port']
            
            config_lines.append(f"    location /fuzz/{uuid}/ {{")
            config_lines.append(f"        proxy_pass http://{ip}:{port}/;")
            config_lines.append("        proxy_set_header Host $http_host;")
            config_lines.append("        proxy_set_header Upgrade $http_upgrade;")
            config_lines.append("        proxy_set_header Connection upgrade;")
            config_lines.append("        proxy_set_header Accept-Encoding gzip;")
            config_lines.append("    }")
            config_lines.append("")
        
        # Add server block end
        config_lines.append("}")
        
        return '\n'.join(config_lines)
    
    def _write_config_file(self, content: str) -> bool:
        """Write the nginx configuration file"""
        try:
            # Ensure the sites-available directory exists
            os.makedirs(NGINX_SITES_AVAILABLE, exist_ok=True)
            
            # Write the configuration file
            with open(self.config_path, 'w', encoding='utf-8') as f:
                f.write(content)
            
            logger.info(f"Configuration file written to {self.config_path}")
            return True
        
        except PermissionError:
            logger.error(f"Permission denied writing to {self.config_path}")
            return False
        except Exception as e:
            logger.error(f"Error writing configuration file: {e}")
            return False
    
    def _enable_site(self) -> bool:
        """Enable the nginx site by creating symlink and reloading nginx"""
        try:
            # Create symlink if it doesn't exist
            if not os.path.exists(self.enabled_path):
                os.symlink(self.config_path, self.enabled_path)
                logger.info(f"Created symlink: {self.enabled_path}")
            
            # Test nginx configuration
            test_result = subprocess.run(['nginx', '-t'], capture_output=True, text=True)
            if test_result.returncode != 0:
                logger.error(f"Nginx configuration test failed: {test_result.stderr}")
                return False
            
            # Reload nginx
            reload_result = subprocess.run(['systemctl', 'reload', 'nginx'], capture_output=True, text=True)
            if reload_result.returncode != 0:
                logger.error(f"Nginx reload failed: {reload_result.stderr}")
                return False
            
            logger.info("Nginx configuration reloaded successfully")
            return True
        
        except Exception as e:
            logger.error(f"Error enabling site: {e}")
            return False
    
    def update_configuration(self, instances: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Main method to update nginx configuration"""
        try:
            # Validate input data
            if not self._validate_json_data(instances):
                return {"success": False, "message": "Invalid JSON data format"}
            
            # Generate new configuration
            config_content = self._generate_nginx_config(instances)
            
            # Write configuration file
            if not self._write_config_file(config_content):
                return {"success": False, "message": "Failed to write configuration file"}
            
            # Enable site and reload nginx
            if not self._enable_site():
                return {"success": False, "message": "Failed to enable site or reload nginx"}
            
            return {
                "success": True, 
                "message": f"Successfully updated configuration with {len(instances)} instances"
            }
        
        except Exception as e:
            logger.error(f"Error updating configuration: {e}")
            return {"success": False, "message": f"Internal error: {str(e)}"}
            
    def append_configuration(self, instances: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Append location blocks to existing nginx configuration"""
        try:
            # Validate input data
            if not self._validate_json_data(instances):
                return {"success": False, "message": "Invalid JSON data format"}
                
            # Generate location blocks for the new instances
            location_blocks = []
            for instance in instances:
                uuid = instance['uuid']
                ip = instance['ip']
                port = instance['port']
                
                location_blocks.append(f"    location /fuzz/{uuid}/ {{")
                location_blocks.append(f"        proxy_pass http://{ip}:{port}/;")
                location_blocks.append("        proxy_set_header Host $http_host;")
                location_blocks.append("        proxy_set_header Upgrade $http_upgrade;")
                location_blocks.append("        proxy_set_header Connection upgrade;")
                location_blocks.append("        proxy_set_header Accept-Encoding gzip;")
                location_blocks.append("    }")
                location_blocks.append("")
            
            # Check if configuration file exists
            if os.path.exists(self.config_path):
                # Read existing configuration
                with open(self.config_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                
                # Insert new location blocks before the closing bracket
                if "}" in content:
                    content = content.rstrip()
                    if content.endswith("}"):
                        content = content[:-1]  # Remove the closing bracket
                        content = content + "\n" + "\n".join(location_blocks) + "\n}"
                    else:
                        # Malformed config, append at the end and add closing bracket
                        content = content + "\n" + "\n".join(location_blocks) + "\n}"
                else:
                    # No closing bracket found, create new config
                    content = self._generate_nginx_config(instances)
            else:
                # File doesn't exist, create new configuration
                content = self._generate_nginx_config(instances)
            
            # Write the updated configuration
            if not self._write_config_file(content):
                return {"success": False, "message": "Failed to write configuration file"}
            
            # Enable site and reload nginx
            if not self._enable_site():
                return {"success": False, "message": "Failed to enable site or reload nginx"}
            
            return {
                "success": True,
                "message": f"Successfully appended {len(instances)} instances to configuration"
            }
            
        except Exception as e:
            logger.error(f"Error appending configuration: {e}")
            return {"success": False, "message": f"Internal error: {str(e)}"}
    
    def delete_configuration(self) -> Dict[str, Any]:
        """Delete the nginx configuration file"""
        try:
            # Check if symlink exists and remove it
            if os.path.exists(self.enabled_path):
                os.remove(self.enabled_path)
                logger.info(f"Removed symlink: {self.enabled_path}")
            
            # Check if configuration file exists and remove it
            if os.path.exists(self.config_path):
                os.remove(self.config_path)
                logger.info(f"Deleted configuration file: {self.config_path}")
                
                # Reload nginx to apply changes
                reload_result = subprocess.run(['systemctl', 'reload', 'nginx'], capture_output=True, text=True)
                if reload_result.returncode != 0:
                    logger.error(f"Nginx reload failed: {reload_result.stderr}")
                    return {"success": False, "message": "Deleted configuration file but failed to reload nginx"}
                
                return {"success": True, "message": "Configuration file deleted successfully"}
            else:
                return {"success": True, "message": "Configuration file does not exist"}
                
        except Exception as e:
            logger.error(f"Error deleting configuration: {e}")
            return {"success": False, "message": f"Internal error: {str(e)}"}

def execute_remote_command(ip, port, command):
    """Execute a command on a remote machine using SSH"""
    try:
        # Create SSH client
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        # Connect to remote machine
        logger.info(f"Connecting to {ip}:{port} via SSH")
        try:
            client.connect(
                hostname=ip, 
                port=port, 
                username='root', 
                password='Root.1234', 
                timeout=10
            )
            
            logger.info(f"Successfully connected to {ip}:{port}")
            
            # Execute command
            logger.info(f"Executing command: {command}")
            stdin, stdout, stderr = client.exec_command(command)
            
            # Wait for command to complete
            exit_status = stdout.channel.recv_exit_status()
            
            # Get command output
            output = stdout.read().decode('utf-8').strip()
            error = stderr.read().decode('utf-8').strip()
            
            logger.info(f"Command exit status: {exit_status}")
            logger.info(f"Command output: {output}")
            
            if error:
                logger.warning(f"Command error: {error}")
            
            if exit_status == 0:
                return {
                    "success": True,
                    "message": f"Command executed successfully: {output}"
                }
            else:
                return {
                    "success": False,
                    "message": f"Command failed with exit status {exit_status}: {error}"
                }
                
        except Exception as e:
            logger.error(f"SSH connection error: {str(e)}")
            return {
                "success": False,
                "message": f"SSH connection error: {str(e)}"
            }
        finally:
            client.close()
            logger.info(f"SSH connection to {ip}:{port} closed")
    
    except Exception as e:
        logger.error(f"Error in execute_remote_command: {e}")
        return {"success": False, "message": f"Error: {str(e)}"}

# Initialize the manager
config_manager = NginxConfigManager()

@app.route('/update-config', methods=['POST'])
def update_config():
    """API endpoint to update nginx configuration"""
    try:
        # Get JSON data from request
        data = request.get_json()
        
        if data is None:
            return jsonify({"success": False, "message": "No JSON data provided"}), 400
        
        if not isinstance(data, list):
            return jsonify({"success": False, "message": "Data must be a list"}), 400
        
        if len(data) == 0:
            return jsonify({"success": False, "message": "Empty list provided"}), 400
        
        # Update configuration
        result = config_manager.update_configuration(data)
        
        if result["success"]:
            return jsonify(result), 200
        else:
            return jsonify(result), 500
    
    except Exception as e:
        logger.error(f"API error: {e}")
        return jsonify({"success": False, "message": f"API error: {str(e)}"}), 500

@app.route('/append-config', methods=['POST'])
def append_config():
    """API endpoint to append nginx location configurations"""
    try:
        # Get JSON data from request
        data = request.get_json()
        
        if data is None:
            return jsonify({"success": False, "message": "No JSON data provided"}), 400
        
        if not isinstance(data, list):
            return jsonify({"success": False, "message": "Data must be a list"}), 400
        
        if len(data) == 0:
            return jsonify({"success": False, "message": "Empty list provided"}), 400
        
        # Append to configuration
        result = config_manager.append_configuration(data)
        
        if result["success"]:
            return jsonify(result), 200
        else:
            return jsonify(result), 500
    
    except Exception as e:
        logger.error(f"API error: {e}")
        return jsonify({"success": False, "message": f"API error: {str(e)}"}), 500

@app.route('/delete-config', methods=['DELETE'])
def delete_config():
    """API endpoint to delete nginx configuration file"""
    try:
        # Delete configuration
        result = config_manager.delete_configuration()
        
        if result["success"]:
            return jsonify(result), 200
        else:
            return jsonify(result), 500
    
    except Exception as e:
        logger.error(f"API error: {e}")
        return jsonify({"success": False, "message": f"API error: {str(e)}"}), 500

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({"status": "healthy", "service": "nginx-config-manager"}), 200

@app.route('/', methods=['GET'])
def root():
    """Root endpoint with usage information"""
    usage_info = {
        "service": "Nginx Configuration Manager",
        "endpoints": {
            "/update-config": {
                "method": "POST",
                "description": "Update nginx configuration (overwrites existing config)",
                "payload": "JSON array with objects containing uuid, ip, port"
            },
            "/append-config": {
                "method": "POST",
                "description": "Append to nginx configuration",
                "payload": "JSON array with objects containing uuid, ip, port"
            },
            "/delete-config": {
                "method": "DELETE",
                "description": "Delete nginx configuration file"
            },
            "/health": {
                "method": "GET", 
                "description": "Health check"
            }
        },
        "example_payload": [
            {"uuid": "abc123", "ip": "127.0.0.1", "port": 8080},
            {"uuid": "def456", "ip": "127.0.0.1", "port": 8081}
        ]
    }
    return jsonify(usage_info), 200

if __name__ == '__main__':
    # Check if running as root (needed for nginx operations)
    if os.geteuid() != 0:
        logger.warning("Script is not running as root. Some operations may fail.")
    
    logger.info("Starting Nginx Configuration Manager...")
    logger.info("Available endpoints:")
    logger.info("  POST /update-config - Update nginx configuration (overwrite)")
    logger.info("  POST /append-config - Append to nginx configuration")
    logger.info("  DELETE /delete-config - Delete nginx configuration file")
    logger.info("  GET /health - Health check")
    logger.info("  GET / - Usage information")
    
    # Start the Flask application
    app.run(host='0.0.0.0', port=5000, debug=False)
